package com.crazymaker.springcloud.ratelimit;

import lombok.extern.slf4j.Slf4j;
import org.junit.Test;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 漏桶算法
 */
@Slf4j
public class LeakBucketLimiter {
    // 上次漏水的时间
    private volatile static long lastOutTime = System.currentTimeMillis();

    // 限制每秒处理2个请求
    private static final int rate = 2;

    // 时间间隔
    private static final long interval = 1000;

    // 桶的容量
    private static final long capacity = 20;

    private static AtomicInteger waterLeakBucket = new AtomicInteger(0);

    public synchronized static boolean tryAcquire(String taskId, int turn) {
        // 判断水桶是否为空
        if (waterLeakBucket.get() == 0) {
            lastOutTime = System.currentTimeMillis();
            waterLeakBucket.incrementAndGet();
            return true;
        }

        long now = System.currentTimeMillis();
        // 判断时间区间
        if (now < lastOutTime + interval) {
            // 判断有没有满
            if (waterLeakBucket.get() < capacity) {
                waterLeakBucket.incrementAndGet();
                return true;
            } else {
                // 限流
                return false;
            }
        }

        // 超过1秒的时间区间，计算该时间内已经处理了多少个请求
        int times = Integer.parseInt(((System.currentTimeMillis() - lastOutTime) / 1000L) * rate + "");
        // 剩余的水容量
        int lastTimes = waterLeakBucket.get() - times;
        // 矫正
        lastTimes = Math.max(0, lastTimes);
        waterLeakBucket.set(lastTimes);
        // 重新更新leakTimeStamp
        lastOutTime = System.currentTimeMillis();

        // 矫正之后再次判断
        if (waterLeakBucket.get() < capacity) {
            waterLeakBucket.incrementAndGet();
            return true;
        } else {
            return false;
        }
    }

    ExecutorService executorService = Executors.newFixedThreadPool(5);

    @Test
    public void testLimit() throws Exception{
        CountDownLatch latch = new CountDownLatch(5);
        AtomicInteger failedCount = new AtomicInteger();
        long start = System.currentTimeMillis();
        for (int i = 0; i < 5; i++) {
            executorService.submit(new Runnable() {
                @Override
                public void run() {
                    try {
                        String taskId = String.valueOf(Thread.currentThread().getId());
                        for (int j = 0; j < 20; j++) {
                            boolean res = LeakBucketLimiter.tryAcquire("taskId" + taskId, 1);
                            if (!res) {
                                // 失败的次数
                                failedCount.incrementAndGet();
                            }
                            Thread.sleep(200);
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    } finally {
                        latch.countDown();
                    }
                }
            });
        }

        latch.await();
        long end = System.currentTimeMillis() - start;
        log.info("失败的次数:{}", failedCount.get());
        log.info("成功的次数:{}", (5 * 20) - failedCount.get());
        log.info("失败的比例:{}", failedCount.get() / 100F);
        log.info("耗时:{}", end / 1000F);
    }
}
